/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2023. All rights reserved.
 *
 * Function : z = x + y
 * This sample is a very basic sample that implements vector add on Ascend plaform.
 */
#include "kernel_operator.h"
#include "math.h"
#include <float.h>
#include<stdio.h>
#include <stdlib.h>
#include <assert.h>
// #include "helper2.h"

// tensor num for each queue
using namespace AscendC;






//冒泡排序



template<typename TYPE_X, typename TYPE_Y> class KernelAdd {
    using T = TYPE_X;
public:
    __aicore__ inline KernelAdd() {}



    __aicore__ float HelperMin(float a, float b)
    {
    if (a<b) return a; return b;
    }

    __aicore__ float HelperMax(float a, float b)
    {
    if (a>b) return a; return b;
    }

    __aicore__ inline void MaxMin(float lhs, float rhs, float& min, float& max) {
    if (lhs >= rhs) {
        min = rhs;
        max = lhs;
    } else {
        min = lhs;
        max = rhs;
    }
    }

    struct BoxInfoPtr {
        float score_;
        int16_t index_;
    };


    struct SelectedIndex {
        int64_t batch_index_ = 0;
        int64_t class_index_ = 0;
        int64_t box_index_ = 0;
    };



    __aicore__ inline bool SuppressByIOU(
    // const float* boxes_data, 
    int32_t box_offset, 
    AscendC::GlobalTensor<float> box_data,
    int64_t box_index1, int64_t box_index2,
                            int64_t center_point_box, float iou_threshold) {
    float x1_min{};
    float y1_min{};
    float x1_max{};
    float y1_max{};
    float x2_min{};
    float y2_min{};
    float x2_max{};
    float y2_max{};
    float intersection_x_min{};
    float intersection_x_max{};
    float intersection_y_min{};
    float intersection_y_max{};

    // const float* box1 = boxes_data + 4 * box_index1;
    // const float* box2 = boxes_data + 4 * box_index2;
    float box1[4] = {box_data.GetValue(box_offset+4 * box_index1), 
                    box_data.GetValue(box_offset+4 * box_index1+1), 
                    box_data.GetValue(box_offset+4 * box_index1+2), 
                    box_data.GetValue(box_offset+4 * box_index1+3)};


    float box2[4] = {box_data.GetValue(box_offset+4 * box_index2), 
                    box_data.GetValue(box_offset+4 * box_index2+1), 
                    box_data.GetValue(box_offset+4 * box_index2+2), 
                    box_data.GetValue(box_offset+4 * box_index2+3)};

    // center_point_box_ only support 0 or 1
    if (0 == center_point_box) {
        // boxes data format [y1, x1, y2, x2],
        MaxMin(box1[1], box1[3], x1_min, x1_max);
        MaxMin(box2[1], box2[3], x2_min, x2_max);

        intersection_x_min = HelperMax(x1_min, x2_min);
        intersection_x_max = HelperMin(x1_max, x2_max);
        if (intersection_x_max <= intersection_x_min)
        return false;

        MaxMin(box1[0], box1[2], y1_min, y1_max);
        MaxMin(box2[0], box2[2], y2_min, y2_max);
        intersection_y_min = HelperMax(y1_min, y2_min);
        intersection_y_max = HelperMin(y1_max, y2_max);
        if (intersection_y_max <= intersection_y_min)
        return false;
    } else {
        // 1 == center_point_box_ => boxes data format [x_center, y_center, width, height]
        float box1_width_half = box1[2] / 2;
        float box1_height_half = box1[3] / 2;
        float box2_width_half = box2[2] / 2;
        float box2_height_half = box2[3] / 2;

        x1_min = box1[0] - box1_width_half;
        x1_max = box1[0] + box1_width_half;
        x2_min = box2[0] - box2_width_half;
        x2_max = box2[0] + box2_width_half;

        intersection_x_min = HelperMax(x1_min, x2_min);
        intersection_x_max = HelperMin(x1_max, x2_max);
        if (intersection_x_max <= intersection_x_min)
        return false;

        y1_min = box1[1] - box1_height_half;
        y1_max = box1[1] + box1_height_half;
        y2_min = box2[1] - box2_height_half;
        y2_max = box2[1] + box2_height_half;

        intersection_y_min = HelperMax(y1_min, y2_min);
        intersection_y_max = HelperMin(y1_max, y2_max);
        if (intersection_y_max <= intersection_y_min)
        return false;
    }

    const float intersection_area = (intersection_x_max - intersection_x_min) *
                                    (intersection_y_max - intersection_y_min);

    if (intersection_area <= .0f) {
        return false;
    }

    const float area1 = (x1_max - x1_min) * (y1_max - y1_min);
    const float area2 = (x2_max - x2_min) * (y2_max - y2_min);
    const float union_area = area1 + area2 - intersection_area;

    if (area1 <= .0f || area2 <= .0f || union_area <= .0f) {
        return false;
    }

    const float intersection_over_union = intersection_area / union_area;

    return intersection_over_union > iou_threshold;
    }



    __aicore__ void MergeArr(int16_t *a,int begin1,int end1,int begin2,int end2,int16_t *tmp, int box_score_offset)
    {
        int left = begin1, right = end2;
        int index = begin1;
        while (begin1 <= end1 && begin2 <= end2)
        {
            auto index_j = a[begin1];
            auto score_j = score.GetValue(a[begin1]+box_score_offset);
                // auto score_j = scores_in_class.GetValue(data[j]);
            auto index_j1 = a[begin2];
            auto score_j1 = score.GetValue(a[begin2]+box_score_offset);

            if (score_j<score_j1||(score_j==score_j1 && index_j>index_j1))
            {
                tmp[index++] = a[begin1++];
            }
            else
            {
                tmp[index++] = a[begin2++];
            }
        }
        while (begin1 <= end1)
            tmp[index++] = a[begin1++];

        while (begin2 <= end2)
            tmp[index++] = a[begin2++];

        //把归并好的再tmp的数据拷贝回原数组
        for (int i = left; i <= right; ++i)
        {
            a[i] = tmp[i];
        }
    }

    __aicore__ void MergeSortNonR(int16_t* a, int n, int16_t* tmp, int box_score_offset)
    {
        int gap = 1;
        while (gap < n)
        {
            for (int i = 0; i < n; i += 2 * gap)
            {
                //控制[i,i+gap-1]  [i+gap,i+2*gap-1];
                int begin1 = i, end1 = i + gap - 1;
                int begin2 = i + gap, end2 = i + 2 * gap - 1;
                //1.合并的时候只有第一组，第二组一个数据都没有(就不需要合并)
                if (begin2>=n)
                {
                    break;
                }

                //2.合并的时候第二组只有部分数据，需要修改end2边界
                if (end2 >= n)
                {
                    end2 = n - 1;
                }
                MergeArr(a, begin1, end1, begin2, end2, tmp, box_score_offset); 
            }
            gap *= 2;
        }
    }

    __aicore__ inline void Init(GM_ADDR box_addr,GM_ADDR score_addr, GM_ADDR y ,
     int16_t* result,int32_t outputlength, 
     int32_t box_input_num, int32_t score_input_num, int32_t output_num, int32_t max_output_boxes_per_class,
     float iou_threshold, float score_threshold, int32_t center_point_box, int32_t num_batches_, 
     int32_t num_classes_, int32_t num_boxes_
     )
    {
        box.SetGlobalBuffer((__gm__ float *)box_addr, box_input_num);
        score.SetGlobalBuffer((__gm__ float *)score_addr, score_input_num);
        yGm.SetGlobalBuffer((__gm__ int32_t *)y, output_num);
        pipe.InitBuffer(tmp1, 4000 * sizeof(float));
        

        auto scores_in_class = tmp1.Get<float>();
        int16_t candidate_boxes[3400];
        int16_t tmp[3400];

        int32_t selected_indices_index = 0;

        int32_t selected_boxes_inside_class_index;

        for (int64_t batch_index = 0; batch_index < num_batches_; ++batch_index) {
            for (int64_t class_index = 0; class_index < num_classes_; ++class_index) {

            int64_t box_score_offset = (batch_index * num_classes_ + class_index) * num_boxes_;
            // const float* batch_boxes = boxes_data + (batch_index * num_boxes_ * 4);
            int32_t box_offset = batch_index * num_boxes_ * 4;

            

            int32_t candidate_boxes_index = 0;
            // candidate_boxes.reserve(num_boxes_);

            // Filter by score_threshold_
            // const auto* class_scores = scores_data + box_score_offset;
            int32_t score_offset = box_score_offset;
            // if (pc.score_threshold_ != nullptr) {
            if (1){
                for (int64_t box_index = 0; box_index < num_boxes_; ++box_index, ++score_offset) {
                if ( score.GetValue(score_offset)> score_threshold) {
                    candidate_boxes[candidate_boxes_index] = int16_t(box_index);
                    // scores_in_class.SetValue(candidate_boxes_index, score.GetValue(score_offset));

                    candidate_boxes_index+=1;
                }
                }
            } 

            // else {
            //     for (int64_t box_index = 0; box_index < num_boxes_; ++box_index, ++class_scores) {
            //     candidate_boxes[candidate_boxes_index].score_ = *class_scores;
            //     candidate_boxes[candidate_boxes_index].index_ = box_index;
            //     candidate_boxes_index+=1;
            //     // candidate_boxes.emplace_back(*class_scores, box_index);
            //     }
            // }

            // std::priority_queue<BoxInfoPtr, std::vector<BoxInfoPtr>> sorted_boxes(std::less<BoxInfoPtr>(), std::move(candidate_boxes));

            


            // bubbleSort(candidate_boxes, candidate_boxes_index, box_score_offset, scores_in_class);
            MergeSortNonR(candidate_boxes, candidate_boxes_index, tmp, box_score_offset);

            selected_boxes_inside_class_index = 0;
            // Get the next box with top score, filter by iou_threshold
            while (candidate_boxes_index>0 && static_cast<int64_t>(selected_boxes_inside_class_index) < max_output_boxes_per_class) {
                // const BoxInfoPtr& next_top_score = candidate_boxes[candidate_boxes_index-1];
                
                auto box_index_to_select = candidate_boxes[candidate_boxes_index-1];

                bool selected = true;

                for (int i=0; i<selected_boxes_inside_class_index; ++i){

                    int32_t current_index = yGm.GetValue((selected_indices_index-1-i)*3+2);

                    if (SuppressByIOU(box_offset, box, box_index_to_select, current_index, center_point_box, iou_threshold)) {
                        selected = false;
                        break;
                    }
                }

                if (selected) {
                selected_boxes_inside_class_index+=1;
                yGm.SetValue(selected_indices_index*3, int32_t(batch_index));
                yGm.SetValue(selected_indices_index*3+1, int32_t(class_index));
                yGm.SetValue(selected_indices_index*3+2, int32_t(box_index_to_select));
                selected_indices_index+=1;
                
                }
                // sorted_boxes.pop();
                candidate_boxes_index-=1;

            }  // while

            
            }  // for class_index
        }  // for batch_index



        // for(int32_t output_index = 0; output_index<outputlength2; ++output_index)
        // {
        //     yGm.SetValue(output_index, int32_t(result2[output_index]));
        // }
    }

private:
    // GlobalTensor<float> xGm;
    AscendC::TPipe pipe;
    GlobalTensor<float> box, score;
    GlobalTensor<int32_t> yGm;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> tmp1;

};

extern "C" __global__ __aicore__ void non_max_suppression(GM_ADDR x1, GM_ADDR x2, GM_ADDR x3, GM_ADDR x4, GM_ADDR x5, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    KernelAdd<float, float> op;
    op.Init(x1, x2, y, tiling_data.result, tiling_data.outputlength,
    tiling_data.box_input_num, tiling_data.score_input_num, tiling_data.output_num, tiling_data.max_output_boxes_per_class, 
    tiling_data.iou_threshold, tiling_data.score_threshold, tiling_data.center_point_box, tiling_data.num_batches_, 
    tiling_data.num_classes_, tiling_data.num_boxes_
    );  
}
